# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
操作符发射器

处理各种操作符的代码生成。
"""

from ...ir import IRNode
from ...ir.expressions import IRBinOp, IRUnaryOp, IRBoolOp, IRCompare, IRConditional
from ..core.emitter import (
    EmitResult,
    OperatorEmitter as BaseOperatorEmitter,
)
from ..core.context import CodeGenContext
from ..configs.operators import get_operator_mapping, get_operator_precedence
from ..core.generator import get_global_generator

# 支持的操作符类型
supported_ops = [
    # 算术操作符
    "+",
    "-",
    "*",
    "/",
    "//",
    "%",
    "**",
    # 比较操作符
    "==",
    "!=",
    "<",
    "<=",
    ">",
    ">=",
    "is",
    "is not",
    # 逻辑操作符
    "and",
    "or",
    "not",
    # 一元操作符
    "neg",  # 负号
    "pos",  # 正号
    "invert",  # 按位取反
    # 位操作符
    "&",
    "|",
    "^",
    "~",
    "<<",
    ">>",
    # 三元运算符
    "?:",
]


class OperatorEmitter(BaseOperatorEmitter):
    """操作符发射器

    处理二元操作、一元操作、布尔操作、比较操作等。
    """

    def __init__(self):
        super().__init__(supported_ops)

    def can_emit(self, node: IRNode, context: CodeGenContext) -> bool:
        """检查是否可以处理该节点"""
        # 处理条件表达式（三元运算符）
        if isinstance(node, IRConditional):
            return True

        # 其他操作符使用父类的逻辑
        return super().can_emit(node, context)

    def emit(self, node: IRNode, context: CodeGenContext) -> EmitResult:
        """生成操作符代码"""
        if isinstance(node, IRBinOp):
            return self._emit_binary_op(node, context)
        elif isinstance(node, IRUnaryOp):
            return self._emit_unary_op(node, context)
        elif isinstance(node, IRBoolOp):
            return self._emit_bool_op(node, context)
        elif isinstance(node, IRCompare):
            return self._emit_compare_op(node, context)
        elif isinstance(node, IRConditional):
            return self._emit_conditional(node, context)
        else:
            return EmitResult(code="/* Unknown operator */", precedence=0)

    def _emit_binary_op(self, node: IRBinOp, context: CodeGenContext) -> EmitResult:
        """生成二元操作符代码"""

        generator = get_global_generator()
        op_precedence = get_operator_precedence(node.op)

        # 生成左右操作数
        left_code = generator.generate_expression(
            node.left, context, parent_precedence=op_precedence
        )
        right_code = generator.generate_expression(
            node.right, context, parent_precedence=op_precedence
        )

        # 获取C++操作符
        cpp_op = get_operator_mapping(node.op)

        # 特殊处理某些操作符
        if node.op == "**":
            # 幂运算需要使用pow函数
            context.require_include("cmath")
            code = f"std::pow({left_code}, {right_code})"
        elif node.op == "//":
            # 整数除法
            code = f"({left_code} / {right_code})"
        else:
            # 普通二元操作符
            code = f"{left_code} {cpp_op} {right_code}"

        return EmitResult(code=code, precedence=op_precedence)

    def _emit_unary_op(self, node: IRUnaryOp, context: CodeGenContext) -> EmitResult:
        """生成一元操作符代码"""

        generator = get_global_generator()
        op_precedence = get_operator_precedence("unary")

        # 生成操作数
        operand_code = generator.generate_expression(
            node.operand, context, parent_precedence=op_precedence
        )

        # 获取C++操作符
        if node.op == "not":
            cpp_op = "!"
        elif node.op == "neg":
            cpp_op = "-"
        elif node.op == "pos":
            cpp_op = "+"
        elif node.op == "invert":
            cpp_op = "~"
        else:
            cpp_op = get_operator_mapping(node.op)

        code = f"{cpp_op}{operand_code}"
        return EmitResult(code=code, precedence=op_precedence)

    def _emit_bool_op(self, node: IRBoolOp, context: CodeGenContext) -> EmitResult:
        """生成布尔操作符代码"""

        generator = get_global_generator()

        # 获取操作符和优先级
        cpp_op = "&&" if node.op == "and" else "||"
        op_precedence = get_operator_precedence(node.op)

        # 生成所有操作数
        parts = []
        for value in node.values:
            # 对于布尔运算符的子表达式，需要特殊处理括号
            if isinstance(value, IRBoolOp):
                # 如果子表达式是不同优先级的布尔运算符，需要加括号
                child_prec = get_operator_precedence(value.op)
                if child_prec != op_precedence:
                    # 使用一个较小的parent_prec来强制加括号
                    part_code = generator.generate_expression(
                        value, context, parent_precedence=child_prec - 1
                    )
                else:
                    part_code = generator.generate_expression(
                        value, context, parent_precedence=op_precedence
                    )
            else:
                # 对于OR运算符，需要为比较运算符和二元运算符添加括号
                if node.op == "or" and isinstance(value, (IRCompare, IRBinOp)):
                    # OR运算符优先级低于比较和二元运算符，需要括号
                    part_code = generator.generate_expression(
                        value, context, parent_precedence=1
                    )
                else:
                    # 对于其他情况，使用正常的优先级处理
                    part_code = generator.generate_expression(
                        value, context, parent_precedence=op_precedence
                    )
            parts.append(part_code)

        code = f" {cpp_op} ".join(parts)
        return EmitResult(code=code, precedence=op_precedence)

    def _emit_compare_op(self, node: IRCompare, context: CodeGenContext) -> EmitResult:
        """生成比较操作符代码"""

        generator = get_global_generator()
        op_precedence = get_operator_precedence("compare")

        # 生成比较链
        parts = []
        prev_code = generator.generate_expression(
            node.left, context, parent_precedence=op_precedence
        )

        for op, comparator in zip(node.ops, node.comparators):
            right_code = generator.generate_expression(
                comparator, context, parent_precedence=op_precedence
            )

            # 特殊处理 in 和 not in 运算符
            if op == "in":
                # item in container -> container.__contains__(item)
                parts.append(f"{right_code}.__contains__({prev_code})")
            elif op == "not in":
                # item not in container -> !container.__contains__(item)
                parts.append(f"!{right_code}.__contains__({prev_code})")
            else:
                # 获取C++比较操作符
                cpp_op = get_operator_mapping(op)
                parts.append(f"{prev_code} {cpp_op} {right_code}")

            prev_code = right_code

        # 用 && 连接多个比较
        code = " && ".join(parts)
        return EmitResult(code=code, precedence=op_precedence)

    def _emit_conditional(
        self, node: IRConditional, context: CodeGenContext
    ) -> EmitResult:
        """生成条件表达式（三元运算符）代码

        生成 C++ 的三元运算符：condition ? true_expr : false_expr

        Args:
            node: IRConditional 节点
            context: 代码生成上下文

        Returns:
            EmitResult: 生成的代码和优先级
        """

        generator = get_global_generator()

        # 三元运算符的优先级（C++ 中为 16，较低）
        conditional_precedence = 16

        # 生成条件表达式
        test_code = generator.generate_expression(
            node.test, context, parent_precedence=conditional_precedence
        )

        # 生成 true 分支表达式
        body_code = generator.generate_expression(
            node.body, context, parent_precedence=conditional_precedence
        )

        # 生成 false 分支表达式
        orelse_code = generator.generate_expression(
            node.orelse, context, parent_precedence=conditional_precedence
        )

        # 生成三元运算符代码
        code = f"({test_code}) ? {body_code} : {orelse_code}"

        return EmitResult(code=code, precedence=conditional_precedence)

    def get_priority(self) -> int:
        """操作符发射器的优先级"""
        return 100  # 中等优先级
